package com.zyk.leetcode;

import java.util.HashMap;
import java.util.Map;

/**
 * @author zhangsan
 * @date 2021/6/6 11:39
 */
public class C474 {

    // 纯暴力解
    public static int findMaxForm(String[] strs, int m, int n) {
        return process(strs, 0, m, n);
    }

    public static int process(String[] strs, int i, int restZero, int restOne) {
        if (i == strs.length) return 0;
        int oneCnt = 0;
        int zeroCnt = 0;
        for (char c : strs[i].toCharArray()) {
            if (c == '0') zeroCnt++;
            else oneCnt++;
        }
        int p1 = 0;
        if (restOne >= oneCnt && restZero >= zeroCnt) {
            p1 = 1 + process(strs, i + 1, restZero - zeroCnt, restOne - oneCnt);         // 要这个数
        }
        int p2 = process(strs, i + 1, restZero, restOne);         // 不要这个数
        return Math.max(p1, p2);
    }


    // 三维动态规划
    public static int dp(String[] strs, int m, int n) {
        int N = strs.length;
        int[] oneCnt = new int[N];
        for (int i = 0; i < strs.length; i++) {
            for (int j = 0; j < strs[i].length(); j++) {
                if (strs[i].charAt(j) == '1') oneCnt[i]++;
            }
        }
        int[][][] dp = new int[N + 1][m + 1][n + 1];
        for (int i = N - 1; i >= 0; --i) {
            for (int j = 0; j <= m; ++j) {
                for (int k = 0; k <= n; ++k) {

                    int p1 = 0;
                    int zeroCnt = strs[i].length() - oneCnt[i];
                    if (j >= zeroCnt && k >= oneCnt[i]) {
                        p1 = 1 + dp[i + 1][j - zeroCnt][k - oneCnt[i]];         // 要这个数
                    }
                    int p2 = dp[i + 1][j][k];         // 不要这个数
                    dp[i][j][k] = Math.max(p1, p2);

                }
            }
        }

        return dp[0][m][n];
    }

    // 记忆化搜索
    public static int findMaxForm2(String[] strs, int m, int n) {
        Map<Integer, Map<Integer, Map<Integer, Integer>>> dp = new HashMap<>();
        for (int i = 0; i < strs.length; i++) {
            dp.put(i, new HashMap<>());
        }
        return process2(strs, 0, m, n, dp);
    }

    public static int process2(String[] strs, int i, int restZero, int restOne, Map<Integer, Map<Integer, Map<Integer, Integer>>> dp) {
        if (i == strs.length) return 0;
        if (dp.get(i).containsKey(restOne) && dp.get(i).get(restOne).containsKey(restZero))
            return dp.get(i).get(restOne).get(restZero);

        int oneCnt = 0;
        int zeroCnt = 0;
        for (char c : strs[i].toCharArray()) {
            if (c == '0') zeroCnt++;
            else oneCnt++;
        }
        int p1 = 0;
        if (restOne >= oneCnt && restZero >= zeroCnt) {
            p1 = 1 + process2(strs, i + 1, restZero - zeroCnt, restOne - oneCnt, dp);         // 要这个数
        }
        int p2 = process2(strs, i + 1, restZero, restOne, dp);         // 不要这个数
        int ans = Math.max(p1, p2);
        if (!dp.get(i).containsKey(restOne)) {
            dp.get(i).put(restOne, new HashMap<>());
        }
        dp.get(i).get(restOne).put(restZero, ans);
        return ans;
    }


    // for test
    public static void main(String[] args) {
//        String[] strs = {"10", "0001", "111001", "1", "0"};
//        int m = 5, n = 3;

        String[] strs = {"10", "0001", "111001", "1", "0"};
        int m = 4, n = 3;
        System.out.println(findMaxForm(strs, m, n));
        System.out.println(findMaxForm2(strs, m, n));
        System.out.println(dp(strs, m, n));
    }

}
